package com.bzframework.pro.authorize.filter;

import com.bzframework.pro.authorize.context.ReactiveUserTokenContextHolder;
import com.bzframework.pro.authorize.token.holder.UserTokenHolder;
import com.bzframework.pro.authorize.token.parser.ParsedToken;
import com.bzframework.pro.authorize.token.parser.TokenParser;
import lombok.RequiredArgsConstructor;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;

import java.util.function.Function;

/**
 * @Author bingzhong
 * @Date 2024/11/23 22:24
 * @Description
 */
@RequiredArgsConstructor
public class ReactiveRequestContextWebFilter implements WebFilter {

    private final TokenParser tokenParser;
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {

        return tokenParser.parseToken(exchange)
                          .map(token -> Mono.fromRunnable(()-> UserTokenHolder.setCurrentToken(token))
                                            .then(chain.filter(exchange)
                                                       .contextWrite(Context.of(ParsedToken.class, token))))
                          .defaultIfEmpty(chain.filter(exchange))
                          .flatMap(Function.identity())
                          .contextWrite(Context.of(ReactiveUserTokenContextHolder.REQUEST_CONTEXT_KEY, exchange.getRequest()));

    }
}
